# -*- coding: utf-8 -*-
"""
@Time: 2021-10-11 13:06
@Author: 鹄望潇湘
@File: SocketLib.py
@desc: 使用Socket进行全双工通信的服务器与客户端，发送数据的单个数据文件长度最大支持4GB。使用时，服务器端继承SocketServer类，
客户端继承SocketClient类，并重写相关方法。该代码支持 Python2.x 与 Python 3.x
关于可发送的数据类型问题。Socket客户端与服务端支持发送字符串和bytes两种类型的数据。如果需要发送混合类型的数据，或者数据类型在不同运行
时刻会发生变换，推荐使用本文件中的Message结构。该结构可传输混合传输文本与bytes类型的数据(如图片，音频，
视频)等。
Message类还提供了将PIL.Image类型转换为bytes和将bytes转换为PIL.Image的static函数

"""
import io
import logging
import socket
import struct
import threading
import sys
import time

from PIL import Image

if sys.version < '3':
    reload(sys)  # 此处代码为python2设计，在python3环境下会标红，但是整个代码将正常运行，无需做额外处理。
    sys.setdefaultencoding('utf8')


class SocketServer:
    """
    使用Socket通信的服务端，message_process(self, message, host)函数必须重写\n
    可选重写函数:\n
    connected(host) 某个客户端连接时的回调函数\n
    disconnected(host) 某个客户端断开连接时的回调函数\n

    """
    def __init__(self):
        self.host = None
        self.port = None
        self.max_connection = 5
        self.socket = None
        self.connected_host = dict()
        self.stop = False
        self.logger = None
        self.init_logger()

    def send_bytes(self, host, data):
        """
        向服务器发送 bytes 类型的数据，json类型的数据可以使用 json_obj.encode()函数将其转换为bytes类型

        :param data: 将要发送的bytes类型的数据
        :return: None
        """
        try:
            connection = self.connected_host[host]
            length = FrameMessage.int_to_bytes(len(data), 4, 'big')
            connection.sendall(length + data)
        except:
            self.logger.info("数据发送失败")
            pass

    def send_string(self, host, data):
        """
        向服务器端发送字符串类型的数据

        :param data: 即将发送的数据
        :return: None
        """
        self.send_bytes(host, data.encode('utf-8'))

    def message_process(self, message, host):
        """
        处理请求发过来的数据。message的类型为bytes。该函数必须重写

        :param message: 数据信息，数据类型为bytes
        :param host: 发送该消息的客户端地址, 类型为str
        :return: None
        """
        return NotImplementedError

    def start(self, ip, port):
        """开启服务器监听，该函数会阻塞当前线程"""
        self.host = ip
        self.port = port
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.bind((self.host, self.port))
        self.socket.listen(self.max_connection)
        self.logger.info("服务器已启动，监听地址："+self.host+":"+str(self.port))
        while self.stop is False:
            try:
                conn, addr = self.socket.accept()
                self.connected_host[str(addr)] = conn
                self.logger.info("已收到来自"+str(addr)+"的连接")
                threading.Thread(target=self.message_handle, args=(conn, str(addr))).start()
            except:
                break
        self.stop = True
        self.logger.info("服务器已关闭")

    def init_logger(self):
        self.logger = logging.getLogger("SocketServer")
        self.logger.setLevel(logging.INFO)
        console_handle = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s : %(levelname)s %(message)s')
        console_handle.setFormatter(formatter)
        self.logger.addHandler(console_handle)

    def message_handle(self, connection, address):
        """
        与某个客户端建立连接后，监听该客户端发过来的消息函数。
        :param connection: 与客户端的socket连接
        :param address: 客户端的IP与端口
        :return: None
        """
        # self.logger.info("开始处理消息")
        self.connected(address)
        buffer = bytes()
        while True:
            data = []
            try:
                data = bytearray(connection.recv(4096))
            except:
                pass
            if len(data) == 0:
                self.logger.info(str(address)+"已断开连接")
                self.disconnected(str(address))
                break
            buffer = buffer + data
            while True:  # 处理一次收到多个数据
                message = FrameMessage(buffer)
                buffer = message.remains
                if len(message.complete_message) is not 0:
                    self.message_process(message.complete_message, address)
                else:
                    break
        pass

    def disconnect(self, host):
        """与某一个特定的客户端断开连接"""
        if host in self.connected_host.keys():
            self.connected_host[host].close()
            self.connected_host.pop(host)
            self.disconnected(host)
        pass

    def closeServer(self):
        """断开所有的连接并关闭服务器"""
        self.stop = True
        for item in self.connected_host.keys():
            self.connected_host[item].close()
        self.connected_host.clear()
        self.socket.close()

    def connected(self, host):
        """
        服务器收到客户端连接后的回调函数
        """
        pass

    def disconnected(self, host):
        """
        客户端断开连接后的回调函数
        """
        pass


class FrameMessage:
    """
    用于分割数据帧，构造函数传入的数据头前两个字节必须是一个数据帧的开始，头2个字节记录着该数据帧所包含的数据长度
    """
    def __init__(self, data):
        self.complete_message = bytes()
        self.remains = data[0:]
        self.frame_head_length = 4
        if len(data) > self.frame_head_length:  # 防止接收到的数据长度不够帧头长度
            self.calculate(data)

    def calculate(self, data):
        # 取出帧头，确定该数据帧的长度
        # expect_length = int.from_bytes(data[0:self.frame_head_length], byteorder='big', signed=False)
        expect_length = self.bytes_to_int(data[0:self.frame_head_length], 4, mode = 'big')
        data = data[self.frame_head_length:]
        if len(data) > expect_length:
            self.complete_message = data[0: expect_length]
            self.remains = data[expect_length:]
        if len(data) == expect_length:
            self.complete_message = data[0: expect_length]
            self.remains = bytes()

    def get_complete_message(self):
        return self.complete_message

    def get_remain_message(self):
        return self.remains

    @staticmethod
    def bytes_to_int(data, length, mode='big'):
        result = 0
        if mode == 'little':
            data.reverse()
        for i in range(0, length):
            result = result | data[i]
            result = result << 8
        return result >> 8

    @staticmethod
    def int_to_bytes(data, length, mode='big'):
        result = bytearray(length)
        for i in range(0, length):
            result[i] = data & 255
            data = data >> 8
        if mode == 'big':
            result.reverse()
        return result


class SocketClient:
    """
    使用Socket与服务器通信的客户端，
    recv_message(self, message)函数必须重写
    connected(self), disconnected(self)函数可选重写
    """

    def __init__(self, reconnect=5):
        """

        :param reconnect: 连接意外关闭后重连的次数
        """
        self.logger = None
        self.init_logger()
        self.host = None
        self.port = None
        self.should_close = False
        self.socket = None
        self.reconnect_time = reconnect

    def connect(self, ip, port):
        self.host = ip
        self.port = port
        if self.connectServer():
            threading.Thread(target=self.message_listen, args=(self.socket,)).start()
            self.connected()
        else:
            self.reconnect()

    def connectServer(self):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.settimeout(5)
        try:
            self.socket.connect((self.host, self.port))
            self.logger.info("已连接上" + self.host + ":" + str(self.port))
        except Exception as e:
            return False
        return True

    def init_logger(self):
        self.logger = logging.getLogger("SocketClient")
        self.logger.setLevel(logging.INFO)
        console_handle = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s : %(levelname)s %(message)s')
        console_handle.setFormatter(formatter)
        self.logger.addHandler(console_handle)

    def send_bytes(self, data):
        """
        向服务器发送 bytes 类型的数据，json类型的数据可以使用 json_obj.encode()函数将其转换为bytes类型

        :param data: 将要发送的bytes类型的数据
        :return: None
        """
        if self.socket is not None:
            length = FrameMessage.int_to_bytes(len(data), 4, 'big')
            # length = int.to_bytes(len(data), 4, 'big')
            # self.logger.info("客户端即将发送"+str(length+data))
            try:
                self.socket.sendall(length + data)
            except Exception as e:
                self.logger.info("发送数据失败")

    def send_string(self, data):
        """
        向服务器端发送字符串类型的数据

        :param data: 即将发送的数据
        :return: None
        """
        self.send_bytes(data.encode('utf-8'))

    def message_listen(self, connection):
        """接收服务器消息的消息循环函数"""
        # self.logger.info("线程" + threading.currentThread().getName() + "已开始")
        buffer = bytes()
        last_empty = int(time.time())
        while self.should_close is False:
            data = bytearray()
            try:
                data = bytearray(connection.recv(4096))
            except:
                pass
            if len(data) == 0:
                now = int(time.time())
                if self.should_close or self.socket is None:
                    break
                if now-last_empty+1 > self.socket.gettimeout():
                    last_empty = now
                    continue

                if self.reconnect() is False:
                    self.logger.info("已与"+self.host+":"+str(self.port)+"断开连接")
                    self.close()
                break
            buffer = buffer + data
            while True:  # 处理一次收到多个数据
                message = FrameMessage(buffer)
                buffer = message.remains
                if len(message.complete_message) is not 0:
                    self.recv_message(message.complete_message)
                else:
                    break
        # self.logger.info("线程"+threading.currentThread().getName()+"已结束")
        pass

    def recv_message(self, message):
        """
        处理从服务器发送过来的消息的函数，该函数必须重写

        :param message: 从服务端发送过来的消息
        :return: None
        """
        return NotImplementedError

    def close(self):
        """
        断开与服务器的socket连接

        :return: None
        """
        self.should_close = True
        try:
            self.socket.close()
        except:
            pass
        self.socket = None
        self.disconnected()

    def connected(self):
        """
        连接上服务器后的回调函数

        :return: None
        """
        pass

    def disconnected(self):
        """
        与服务器断开连接后的回调函数

        :return: None
        """
        pass

    def reconnect(self):
        if self.reconnect_time == 0:
            return False
        sleep_time = 5
        for i in range(0, self.reconnect_time):

            self.logger.info("正在尝试第"+str(i)+"次重连")
            if self.connectServer():
                self.should_close = False
                threading.Thread(target=self.message_listen, args=(self.socket,)).start()
                self.logger.info("已经与服务器重新连接")
                self.connected()
                return True
            else:
                time.sleep(sleep_time)
                sleep_time = sleep_time * 2
                continue
        self.logger.error("经过多次尝试，服务器重连仍然失败")
        return False


class Message:
    """
    将各种常用的数据类型转换为bytes类型以及从bytes类型转换回基本数据类型。
    支持的数据类型为 int, str, float list, dict, bytes，支持上述各种数据类型的混合。
    使用时,直接使用Msg.encode(data)函数将数据编码为bytes，使用Msg.decode(data)将Bytes数据转换回原始数据。
    类部数据类型代码对应关系：
    1 int
    2 str
    3 float
    4 list
    5 dict
    6 bytes
    """
    TYPE_INT = 1
    TYPE_STR = 2
    TYPE_FLOAT = 3
    TYPE_LIST = 4
    TYPE_DICT = 5
    TYPE_BYTES = 6

    @staticmethod
    def encode(data):
        """
        将数据转换为bytes

        :param data: 需要转换的数据支持的数据类型为 int, float, str, bytes, list, dict
        :return: Bytes 转换后的bytes类型的数据
        """
        data_type = type(data)
        if data_type is type(1):
            return struct.pack("biq", Message.TYPE_INT, 8, data)

        if data_type is type('strs') or data_type is type(u'unicode'):
            try:
                byte_str = data.encode('utf-8')
                byte_head = struct.pack("bi", Message.TYPE_STR, len(byte_str))
            except:
                byte_head = struct.pack("bi", Message.TYPE_BYTES, len(data))
                byte_str = data
            return byte_head + byte_str

        if data_type is type(3.145):
            return struct.pack("bid", Message.TYPE_FLOAT, 8, data)
            pass

        if data_type is type([]):
            length = 0
            data_byte = bytearray()
            for item in data:
                item_byte = Message.encode(item)
                length = length + len(item_byte)
                data_byte = data_byte + item_byte
            head = struct.pack("bi", Message.TYPE_LIST, length)
            return head + data_byte
            pass

        if data_type is type(dict()):
            length = 0
            data_byte = bytearray()
            for key in data.keys():
                key_byte = Message.encode(key)
                value_byte = Message.encode(data[key])
                length = length + len(key_byte) + len(value_byte)
                data_byte = data_byte + key_byte + value_byte
            head  = struct.pack("bi", Message.TYPE_DICT, length)
            return head+data_byte

        if data_type is type(b'1') or data_type is type(bytearray()):
            data_length = len(data)
            byte_head = struct.pack("bi", Message.TYPE_BYTES, data_length)
            return byte_head + data

        raise TypeError("不支持的编码数据类型", data_type)

    @staticmethod
    def decode(data):
        """
        将Bytes类型的数据还原

        :param data: io.BytesIO类型，从bytes转换为io.BytesIO请用 data = io.BytesIO(data)
        :return: 还原后的数据，可能为 int, str, float, bytes, dict, list
        """
        if type(data) is type(b'1') or type(data) is type(bytearray()):
            data = io.BytesIO(data)
        head_bytes = data.read(8)
        if len(head_bytes) == 0:
            return None
        data_type, length = struct.unpack('bi', head_bytes)
        extra_bytes = data.read(length)
        if data_type is 1:
            return struct.unpack("q", extra_bytes)[0]

        if data_type is 2:
            if sys.version < '3':
                return str.decode(extra_bytes, encoding='utf-8')
            else:
                return extra_bytes.decode(encoding='utf-8')

        if data_type is 3:
            return struct.unpack("bid", head_bytes+extra_bytes)[2]

        if data_type is 4:
            array_data = io.BytesIO(extra_bytes)
            result = []
            while True:
                son_obj = Message.decode(array_data)
                if son_obj is None:
                    break
                else:
                    result.append(son_obj)
            return result

        if data_type is 5:
            array_data = io.BytesIO(extra_bytes)
            result = dict()
            while True:
                key = Message.decode(array_data)
                if key is None:
                    break
                value = Message.decode(array_data)
                result[key] = value
            return result

        if data_type is 6:
            return extra_bytes

    @staticmethod
    def ImagetoBytes(image_data):
        imgByteArr = io.BytesIO()
        image_data.save(imgByteArr, format='JPEG')
        return imgByteArr.getvalue()

    @staticmethod
    def BytestoImage(bytes_data):
        """

        :param bytes_data: 收到的bytes类型的图片数据
        :return: PIL.Image
        """
        bytes_stream = io.BytesIO(bytes_data)
        return Image.open(bytes_stream)
